#!/bin/bash

# =============================================================================
# Dify 智能重启脚本 v1.0 (2025增强版)
# 支持快速重启、完整重启、单服务重启等多种模式
# 增强功能：智能健康检查 + 详细日志记录 + 错误处理 + 状态验证
# =============================================================================

set -euo pipefail
IFS=$'\n\t'

# --- 全局配置 ---
SCRIPT_VERSION="1.0.0"
DEPLOY_DIR="/root/dify"
DATA_DIR="/root/dify_data"
LOG_FILE="/var/log/dify-restart.log"
GITEE_REPO="git@gitee.com:dd0129/dify_test.git"

# --- 颜色输出 ---
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'
EMOJI_SUCCESS="✅"
EMOJI_ERROR="❌"
EMOJI_WARNING="⚠️"
EMOJI_INFO="ℹ️"

# --- 日志函数 ---
log_info() {
    local msg="[INFO] $(date '+%Y-%m-%d %H:%M:%S') $1"
    echo -e "${BLUE}${EMOJI_INFO} ${msg}${NC}"
    echo "$msg" >> "$LOG_FILE"
}

log_success() {
    local msg="[SUCCESS] $(date '+%Y-%m-%d %H:%M:%S') $1"
    echo -e "${GREEN}${EMOJI_SUCCESS} ${msg}${NC}"
    echo "$msg" >> "$LOG_FILE"
}

log_warning() {
    local msg="[WARNING] $(date '+%Y-%m-%d %H:%M:%S') $1"
    echo -e "${YELLOW}${EMOJI_WARNING} ${msg}${NC}"
    echo "$msg" >> "$LOG_FILE"
}

log_error() {
    local msg="[ERROR] $(date '+%Y-%m-%d %H:%M:%S') $1"
    echo -e "${RED}${EMOJI_ERROR} ${msg}${NC}"
    echo "$msg" >> "$LOG_FILE"
}

log_header() {
    echo
    echo -e "${BLUE}======================================${NC}"
    echo -e "${BLUE}  $1${NC}"
    echo -e "${BLUE}======================================${NC}"
    echo
}

# --- 错误处理 ---
trap 'handle_error $LINENO' ERR
handle_error() {
    log_error "脚本在第 $1 行执行失败"
    log_error "请检查日志文件: $LOG_FILE"
    exit 1
}

# --- 帮助信息 ---
show_help() {
    cat << EOF
Dify 智能重启脚本 v$SCRIPT_VERSION

用法: $0 [选项] [服务名]

选项:
  -h, --help          显示此帮助信息
  -v, --version       显示版本信息
  -q, --quick         快速重启（仅重启容器，不更新代码）
  -f, --full          完整重启（更新代码 + 重启容器）
  -s, --service NAME  重启指定服务
  -c, --check         仅检查服务状态，不执行重启
  -l, --logs [SERVICE] 查看服务日志
  --no-pull          跳过代码更新（即使在完整模式下）
  --force            强制重启，跳过健康检查

重启模式:
  快速模式: 仅重启Docker容器，保持当前代码版本
  完整模式: 拉取最新代码，然后重启所有容器
  服务模式: 仅重启指定的单个服务

服务列表:
  api, web, worker, nginx, postgres, redis, weaviate, sandbox, ssrf_proxy

示例:
  $0                    # 默认快速重启
  $0 -f                 # 完整重启（更新代码）
  $0 -s api             # 仅重启API服务
  $0 -c                 # 检查服务状态
  $0 -l api             # 查看API服务日志
  $0 --force -f         # 强制完整重启

EOF
}

# --- 检查环境 ---
check_environment() {
    log_info "检查运行环境..."
    
    # 检查root权限
    if [[ $EUID -ne 0 ]]; then
        log_error "需要root权限运行此脚本"
        exit 1
    fi
    
    # 检查部署目录
    if [[ ! -d "$DEPLOY_DIR" ]]; then
        log_error "部署目录不存在: $DEPLOY_DIR"
        exit 1
    fi
    
    # 检查Docker和Docker Compose
    if ! command -v docker &> /dev/null; then
        log_error "Docker未安装或不可用"
        exit 1
    fi
    
    if ! docker compose version &> /dev/null; then
        log_error "Docker Compose不可用"
        exit 1
    fi
    
    # 创建日志文件目录
    mkdir -p "$(dirname "$LOG_FILE")"
    touch "$LOG_FILE"
    
    log_success "环境检查通过"
}

# --- 获取服务状态 ---
get_service_status() {
    cd "$DEPLOY_DIR"
    
    if ! docker compose -f docker-compose.production.yml ps &> /dev/null; then
        echo "UNKNOWN"
        return 1
    fi
    
    local running_count=$(docker compose -f docker-compose.production.yml ps --format json 2>/dev/null | jq -r '. | select(.State == "running") | .Service' | wc -l)
    local total_count=$(docker compose -f docker-compose.production.yml ps --format json 2>/dev/null | jq -r '.Service' | wc -l)
    
    echo "$running_count/$total_count"
}

# --- 检查服务健康状态 ---
check_service_health() {
    local service="$1"
    cd "$DEPLOY_DIR"
    
    # 检查容器是否运行
    if ! docker compose -f docker-compose.production.yml ps --format json | jq -r '. | select(.Service == "'$service'") | .State' | grep -q running; then
        echo "STOPPED"
        return 1
    fi
    
    # 检查健康状态
    local health=$(docker compose -f docker-compose.production.yml ps --format json | jq -r '. | select(.Service == "'$service'") | .Health // "unknown"')
    
    case "$health" in
        "healthy") echo "HEALTHY" ;;
        "unhealthy") echo "UNHEALTHY" ;;
        "starting") echo "STARTING" ;;
        *) echo "RUNNING" ;;
    esac
}

# --- 等待服务健康 ---
wait_for_service_health() {
    local service="$1"
    local timeout="${2:-60}"
    local count=0
    
    log_info "等待 $service 服务健康状态..."
    
    while [[ $count -lt $timeout ]]; do
        local status=$(check_service_health "$service")
        
        case "$status" in
            "HEALTHY"|"RUNNING")
                log_success "$service 服务健康状态正常"
                return 0
                ;;
            "STARTING")
                log_info "$service 服务正在启动中... ($count/$timeout)"
                ;;
            "UNHEALTHY")
                log_warning "$service 服务健康检查失败 ($count/$timeout)"
                ;;
            "STOPPED")
                log_error "$service 服务已停止"
                return 1
                ;;
        esac
        
        sleep 2
        ((count += 2))
    done
    
    log_warning "$service 服务在 $timeout 秒内未达到健康状态"
    return 1
}

# --- 检查所有服务状态 ---
check_all_services() {
    log_header "检查服务状态"
    
    cd "$DEPLOY_DIR"
    
    if ! docker compose -f docker-compose.production.yml ps &> /dev/null; then
        log_error "无法获取服务状态，请检查Docker Compose配置"
        return 1
    fi
    
    echo "=== 服务状态详情 ==="
    docker compose -f docker-compose.production.yml ps
    
    echo ""
    echo "=== 服务健康状态 ==="
    
    local services=("api" "web" "worker" "nginx" "postgres" "redis" "weaviate" "sandbox" "ssrf_proxy")
    local healthy_count=0
    local total_count=${#services[@]}
    
    for service in "${services[@]}"; do
        local status=$(check_service_health "$service")
        
        case "$status" in
            "HEALTHY"|"RUNNING")
                log_success "$service: $status"
                ((healthy_count++))
                ;;
            "STARTING")
                log_info "$service: $status"
                ;;
            "UNHEALTHY")
                log_warning "$service: $status"
                ;;
            "STOPPED")
                log_error "$service: $status"
                ;;
        esac
    done
    
    echo ""
    log_info "服务状态汇总: $healthy_count/$total_count 健康"
    
    # 测试外部访问
    echo ""
    echo "=== 外部访问测试 ==="
    
    if curl -f -s http://localhost > /dev/null 2>&1; then
        log_success "Web服务外部访问正常: http://localhost"
    else
        log_warning "Web服务外部访问失败"
    fi
    
    if curl -f -s http://localhost/nginx-health > /dev/null 2>&1; then
        log_success "Nginx健康检查正常: http://localhost/nginx-health"
    else
        log_warning "Nginx健康检查失败"
    fi
}

# --- 安全更新代码 ---
update_code_safely() {
    log_info "安全更新代码..."
    
    cd "$DEPLOY_DIR"
    
    # 保存当前分支和提交
    local current_branch=$(git branch --show-current 2>/dev/null || echo "main")
    local current_commit=$(git rev-parse HEAD 2>/dev/null || echo "unknown")
    
    log_info "当前分支: $current_branch"
    log_info "当前提交: $current_commit"
    
    # 检查工作区状态
    if ! git diff --quiet 2>/dev/null; then
        log_warning "工作区有未提交的更改，将被重置"
        git stash push -m "Auto-stash before restart $(date)" || true
    fi
    
    # 获取远程更新
    log_info "获取远程更新..."
    if ! git fetch origin; then
        log_error "获取远程更新失败，请检查网络连接"
        return 1
    fi
    
    # 检查是否有新提交
    local remote_commit=$(git rev-parse "origin/$current_branch" 2>/dev/null || git rev-parse "origin/main" 2>/dev/null)
    
    if [[ "$current_commit" == "$remote_commit" ]]; then
        log_info "代码已是最新版本，无需更新"
        return 0
    fi
    
    log_info "发现新提交，准备更新..."
    log_info "远程提交: $remote_commit"
    
    # 安全重置到远程版本
    if ! git reset --hard "origin/$current_branch" 2>/dev/null && ! git reset --hard "origin/main" 2>/dev/null; then
        log_error "代码更新失败"
        return 1
    fi
    
    local new_commit=$(git rev-parse HEAD 2>/dev/null)
    log_success "代码更新成功: $current_commit -> $new_commit"
    
    return 0
}

# --- 重启指定服务 ---
restart_service() {
    local service="$1"
    local force="${2:-false}"
    
    cd "$DEPLOY_DIR"
    
    log_info "重启 $service 服务..."
    
    # 检查服务是否存在
    if ! docker compose -f docker-compose.production.yml config --services | grep -q "^$service$"; then
        log_error "服务 $service 不存在"
        return 1
    fi
    
    # 健康检查（除非强制模式）
    if [[ "$force" != "true" ]]; then
        local status=$(check_service_health "$service")
        if [[ "$status" == "HEALTHY" || "$status" == "RUNNING" ]]; then
            log_info "$service 服务当前状态正常 ($status)"
            read -p "是否仍要重启？[y/N] " -n 1 -r
            echo
            if [[ ! $REPLY =~ ^[Yy]$ ]]; then
                log_info "取消重启 $service 服务"
                return 0
            fi
        fi
    fi
    
    # 执行重启
    if ! docker compose -f docker-compose.production.yml restart "$service"; then
        log_error "$service 服务重启失败"
        return 1
    fi
    
    # 等待服务健康
    if ! wait_for_service_health "$service" 30; then
        log_warning "$service 服务重启后健康检查失败"
        return 1
    fi
    
    log_success "$service 服务重启成功"
    return 0
}

# --- 快速重启（仅重启容器） ---
quick_restart() {
    local force="${1:-false}"
    
    log_header "执行快速重启"
    
    cd "$DEPLOY_DIR"
    
    # 记录重启前状态
    local before_status=$(get_service_status)
    log_info "重启前服务状态: $before_status"
    
    # 健康检查（除非强制模式）
    if [[ "$force" != "true" ]]; then
        log_info "执行重启前健康检查..."
        check_all_services
        
        read -p "确认执行快速重启？[y/N] " -n 1 -r
        echo
        if [[ ! $REPLY =~ ^[Yy]$ ]]; then
            log_info "取消重启操作"
            exit 0
        fi
    fi
    
    # 执行重启
    log_info "停止所有服务..."
    if ! docker compose -f docker-compose.production.yml down; then
        log_error "停止服务失败"
        return 1
    fi
    
    log_info "启动所有服务..."
    if ! docker compose -f docker-compose.production.yml up -d; then
        log_error "启动服务失败"
        return 1
    fi
    
    # 等待核心服务健康
    log_info "等待核心服务启动..."
    sleep 10
    
    local core_services=("postgres" "redis" "api" "web" "nginx")
    for service in "${core_services[@]}"; do
        wait_for_service_health "$service" 60 || log_warning "$service 健康检查超时"
    done
    
    # 检查最终状态
    sleep 5
    local after_status=$(get_service_status)
    log_info "重启后服务状态: $after_status"
    
    check_all_services
    
    log_success "快速重启完成"
}

# --- 完整重启（更新代码 + 重启容器） ---
full_restart() {
    local force="${1:-false}"
    local skip_pull="${2:-false}"
    
    log_header "执行完整重启"
    
    # 记录重启前状态
    local before_status=$(get_service_status)
    log_info "重启前服务状态: $before_status"
    
    # 健康检查（除非强制模式）
    if [[ "$force" != "true" ]]; then
        log_info "执行重启前健康检查..."
        check_all_services
        
        read -p "确认执行完整重启（包括代码更新）？[y/N] " -n 1 -r
        echo
        if [[ ! $REPLY =~ ^[Yy]$ ]]; then
            log_info "取消重启操作"
            exit 0
        fi
    fi
    
    # 更新代码
    if [[ "$skip_pull" != "true" ]]; then
        if ! update_code_safely; then
            log_error "代码更新失败，是否继续重启？"
            read -p "继续？[y/N] " -n 1 -r
            echo
            if [[ ! $REPLY =~ ^[Yy]$ ]]; then
                exit 1
            fi
        fi
    else
        log_info "跳过代码更新"
    fi
    
    # 执行重启
    quick_restart "true"
    
    log_success "完整重启完成"
}

# --- 查看服务日志 ---
view_logs() {
    local service="${1:-}"
    
    cd "$DEPLOY_DIR"
    
    if [[ -n "$service" ]]; then
        log_info "查看 $service 服务日志..."
        docker compose -f docker-compose.production.yml logs --tail=50 -f "$service"
    else
        log_info "查看所有服务日志..."
        docker compose -f docker-compose.production.yml logs --tail=20 -f
    fi
}

# --- 主函数 ---
main() {
    # 初始化日志
    log_header "Dify 智能重启脚本 v$SCRIPT_VERSION"
    log_info "脚本启动时间: $(date)"
    log_info "运行参数: $*"
    
    # 解析参数
    local mode="quick"
    local force="false"
    local skip_pull="false"
    local service=""
    local check_only="false"
    local show_logs="false"
    local logs_service=""
    
    while [[ $# -gt 0 ]]; do
        case $1 in
            -h|--help)
                show_help
                exit 0
                ;;
            -v|--version)
                echo "Dify 智能重启脚本 v$SCRIPT_VERSION"
                exit 0
                ;;
            -q|--quick)
                mode="quick"
                shift
                ;;
            -f|--full)
                mode="full"
                shift
                ;;
            -s|--service)
                mode="service"
                service="$2"
                shift 2
                ;;
            -c|--check)
                check_only="true"
                shift
                ;;
            -l|--logs)
                show_logs="true"
                if [[ $# -gt 1 && ! "$2" =~ ^- ]]; then
                    logs_service="$2"
                    shift 2
                else
                    shift
                fi
                ;;
            --no-pull)
                skip_pull="true"
                shift
                ;;
            --force)
                force="true"
                shift
                ;;
            *)
                log_error "未知参数: $1"
                show_help
                exit 1
                ;;
        esac
    done
    
    # 检查环境
    check_environment
    
    # 执行相应操作
    if [[ "$show_logs" == "true" ]]; then
        view_logs "$logs_service"
    elif [[ "$check_only" == "true" ]]; then
        check_all_services
    elif [[ "$mode" == "service" ]]; then
        if [[ -z "$service" ]]; then
            log_error "服务模式需要指定服务名"
            show_help
            exit 1
        fi
        restart_service "$service" "$force"
    elif [[ "$mode" == "full" ]]; then
        full_restart "$force" "$skip_pull"
    else
        quick_restart "$force"
    fi
    
    log_info "脚本执行完成: $(date)"
}

# --- 脚本入口 ---
if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
    main "$@"
fi